import random

import matplotlib.pyplot as plt
import networkx as nx
import pandas as pd
import numpy as np

elseNum=2000
sliceNum=7
effIncrease=1.05

df=pd.read_csv("D:\\personal_file\\2021MCMProblemC_DataSet.CSV")
locDF=df[df["Lab Status"]=="Positive ID"][["Submission Date","Longitude","Latitude"]]
print(locDF)

arr=np.array(locDF[["Longitude","Latitude"]][:sliceNum])

avgArr=np.mean(arr,axis=0)

elsePoints=np.zeros((elseNum,2))
for i in range(elseNum):
    elsePoints[i]=avgArr+np.random.uniform(-1.5,1.5,(2,))

"""predDf=pd.read_csv("D:\\personal_file\\points200.CSV")
elsePoints=np.array(predDf[["1","0"]])"""

allpoints = np.concatenate((arr,elsePoints),axis=0)
allpoints=np.around(allpoints,2)

allnodes= list()
for point in allpoints:
    allnodes.append(tuple(point))

exitnodes=allnodes[:sliceNum]


npos=dict(zip(allnodes,allpoints))


max_iter_num = 150  # 模拟的次数
G = nx.Graph()

for i in range(elseNum+sliceNum):
    G.add_node(allnodes[i],state=0)

for i in range(elseNum+sliceNum):
    for j in range(i+1,elseNum+sliceNum):
        sum=pow(pow(allnodes[i][0]-allnodes[j][0],2)+pow(allnodes[i][1]-allnodes[j][1],2),0.5)
        if (i < sliceNum and j<sliceNum):
            G.add_edge(allnodes[i], allnodes[j], weight=100)
            continue
        if(sum<0.035):
            G.add_edge(allnodes[i], allnodes[j], weight=random.uniform(0, 100))

print(G.edges)

for seed in exitnodes:
    G.node[seed]['state'] = 1  # 表示33是感染的

all_infect_nodes = []  # 所有被感染的节点放在这里
effectList=[]
for seed in exitnodes:
    all_infect_nodes.append(seed)
    effectList.append(1)

infected_graph = nx.DiGraph()  # 被激活的图
for seed in exitnodes:
    infected_graph.add_node(seed)

"""for i in range(14):
    for j in range(i+1,14):
        infected_graph.add_edge(allnodes[i], allnodes[j], weight=100)"""

elseList=[]


for i in range(max_iter_num):
    new_infect = list()  # 新被感染的
    t1 = '%s time' % i + ' %s nodes' % len(all_infect_nodes)
    print(t1)  # 当前有多少个节点被感染

    # 画图
    if(i%5==0):
        fm = pd.DataFrame(elseList)
        fm.to_csv("D:\\personal_file\\points{}天后.CSV".format(i))
        plt.figure(figsize=(15, 7))
        plt.title(t1)
        nx.draw_networkx_nodes(infected_graph,npos,nodelist=exitnodes,node_size=60,node_color='r')
        nx.draw_networkx_nodes(infected_graph, npos, nodelist=elseList, node_size=30,node_color='b')
        #nx.draw_networkx_labels(infected_graph,npos)
        #nx.draw_networkx_edges(infected_graph,npos,arrowsize=15)
        plt.savefig("D:\\personal_file\\pic{}天后.png".format(i))
        plt.show()

    # 感染的机会不止一次
    for i,v in zip(range(len(effectList)),all_infect_nodes):
        for nbr in G.neighbors(v):
            if G.node[nbr]['state'] == 0:  # 如果这个邻居节点没被感染
                edge_data = G.get_edge_data(v, nbr)
                if (random.uniform(0, 100)*effectList[i]>80):
                    G.node[nbr]['state'] = 1
                    new_infect.append(nbr)
                    infected_graph.add_edge(v, nbr)  # 画图 添加边
                    elseList.append(nbr)
        effectList[i]=min(effectList[i]*effIncrease,2)


    all_infect_nodes.extend(new_infect)  # 将新感染的添加到
    print('all_active_nodes:', all_infect_nodes)




